[MOD-16688] x86 distance kernels: multi-accumulator unrolling + SIMD for smaller dimensions#984
[MOD-16688] x86 distance kernels: multi-accumulator unrolling + SIMD for smaller dimensions#984GuyAv46 wants to merge 5 commits into
Conversation
Break the single-accumulator dependency chain in the x86 distance kernels by accumulating into 2-4 independent SIMD registers (following the pattern already used by the SQ8_FP16 kernels and the ARM implementations), and convert the main loops from do-while to while so the kernels are correct for dimensions smaller than one full chunk. Kernels updated (IP + L2, plus SQ8 Cosine/L2 wrappers that inherit): - FP32/FP64: SSE, AVX, AVX512F - FP16: F16C, AVX512F - SQ8_FP32: SSE4, AVX2, AVX2_FMA, AVX512F_BW_VL_VNNI Dispatcher thresholds lowered where the kernels' residual handling is safe below one full chunk (residual loads never touch memory past the minimum dimension): - FP32: dim >= 8 (was 16), FP64: dim >= 4 (was 8) - FP16: per-tier - F16C >= 8, AVX512F >= 16, AVX512FP16_VL stays 32 - SQ8_FP32: dim >= 8 (was 16) Measured on Intel Core Ultra 7 255U (AVX2/F16C tiers, google-benchmark, 9 repetitions, interleaved, median): - FP16 F16C L2/IP: up to 1.9x at dim 1024 - SQ8_FP32 AVX2_FMA IP: up to 1.8x - FP32 SSE L2: up to 1.45x, FP64 SSE up to 1.7x - FP32/FP64 AVX: 1.1-1.3x - Newly-SIMD small dims: FP32 dim 8-15 1.4-2.8x vs scalar, FP16 dim 8-31 5.7-9.4x vs scalar AVX-512 tiers follow the same structure but were not benchmarked locally (no AVX-512 hardware); loop/tail coverage verified by simulation for all dims up to 4096. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
🛡️ Jit Security Scan Results
✅ No security findings were detected in this PR
Security scan by Jit
|
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #984 +/- ##
==========================================
- Coverage 97.12% 97.12% -0.01%
==========================================
Files 141 141
Lines 8179 8245 +66
==========================================
+ Hits 7944 8008 +64
- Misses 235 237 +2 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
…dulus Match each kernel's dispatch modulus (CHOOSE_IMPLEMENTATION chunk) to its main-loop stride, so all residual-based conditions - including the leftover full blocks that were previously handled with runtime tail checks after the main loop - are hard-coded through the residual template parameter: - AVX FP32 / AVX2 SQ8_FP32(+FMA): dispatch mod 32 (4x8 per iteration) - AVX FP64: dispatch mod 16 (4x4) - AVX512F FP32 / FP16 / VNNI SQ8_FP32: dispatch mod 64 (4x16) - AVX512F FP64: dispatch mod 32 (4x8) The main loops are now branch-free except for the loop condition itself, and the residual (0..stride-1) is fully unrolled at compile time. Loop/tail coverage re-verified by simulation for all dims up to 4096. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
|
Update (commit 1afb939): all residual-based conditions are now resolved at compile time. Each kernel's dispatch modulus (
SSE / F16C kernels already had stride == modulus, so they were unrolled through the template from the start. Added the 🤖 Generated with Claude Code |
Benchmarking (local 1/2/4-acc sweep on AVX2 hardware + a 2-vs-4 CI run on the
Xeon 8488C) showed the multi-accumulator win shrinks as SIMD width grows:
- non-512 FP32/FP16/SQ8 (SSE/AVX/F16C/AVX2): 4 acc beats 2 by 10-45% -> kept
- FP64 at every width: 4 acc gives nothing over 2 (too few lanes) -> now 2
- all AVX-512 (FP32/FP64/FP16/SQ8_FP32 VNNI): 4 acc at best parity, the two
512-bit FMA units are already saturated by 2 accs at these dims -> now 2
Reverts those kernels to 2 accumulators and narrows each dispatch modulus to
2x the SIMD chunk accordingly (fewer template instantiations, simpler tail):
AVX-512 FP32/FP16/SQ8_FP32 64 -> 32
AVX-512 FP64 32 -> 16
AVX FP64 16 -> 8
SSE FP64 8 -> 4 (adds CASES4 to implementation_chooser)
All residual handling stays compile-time (if constexpr). 1267 test_spaces pass;
AVX-512 kernels compile here and passed basic+sanitizer CI on the 2-acc branch.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Follow-up: tuned accumulator counts per SIMD widthAfter the initial 4-accumulator pass, I benchmarked 1 vs 2 vs 4 accumulators to check where the 4th/3rd accumulator actually pays off:
The multi-accumulator benefit shrinks as SIMD width grows:
Commit 8550b86 drops those kernels to 2 accumulators and narrows each dispatch modulus to 2× the SIMD chunk (AVX-512 FP32/FP16/SQ8 64→32, AVX-512 FP64 32→16, AVX FP64 16→8, SSE FP64 8→4; added |
GuyAv46
changed the title
GuyAv46
changed the title
…lators-small-dims # Conflicts: # src/VecSim/spaces/L2/L2_SSE_FP32.h
- implementation_chooser.h: document chunk=4 as a valid size (CASES4 was added for the SSE FP64 mod-4 dispatch). - implementation_chooser_cleanup.h: #undef CASES4 alongside the other CASES macros. - test_spaces.cpp: the AVX/AVX512F FP32, AVX512F FP64, and AVX2/FMA/VNNI SQ8_FP32 dispatch moduli widened (to 32/16/32), but the optimization-test Range()s still used the old base, leaving residual instantiations for dims past the old span uncovered. Update each Range to `widest_modulus * 2 + 1` (matching the FP16 line): FP32 8..64, FP64 4..32, SQ8_FP32 8..64. 1460 test_spaces pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
2 participants
Jira: MOD-16688
What
Two related optimizations to the x86 distance kernels, porting the lessons learned from the ARM implementations (and from the newer SQ8↔FP16 kernels, which already do both):
add/fmaddwaited for the previous one (3-5 cycle latency chain, while the CPU can issue 2/cycle). The kernels now accumulate into 2-4 independent registers and combine them once before the final reduction — same pattern asIP_AVX512F_SQ8_FP16.h/IP_AVX2_SQ8_FP16.h.dim < 16/32early-outs to the scalar implementation were needed because the main loops weredo-while(always executing one full chunk) — not because the residual loads were unsafe. The main loops are nowwhileloops that can run zero times, and each dispatcher threshold is lowered to the actual minimum implied by its residual handling (the widest full-width load a kernel performs before there is provably enough data).Kernels changed (IP + L2; SQ8 Cosine/L2 wrappers inherit)
The accumulator count is tuned per SIMD width — a 1-vs-2-vs-4 study (see the follow-up comment) showed the benefit shrinks as the vector widens, so wider tiers keep fewer accumulators:
Rationale for the accumulator split: 4 accumulators beat 2 by 10–45% on the narrow FP32/FP16/SQ8 tiers (SSE/AVX/F16C/AVX2), but gave nothing on FP64 (too few lanes) or on any AVX-512 kernel (the two 512-bit FMA units are already saturated by 2 accumulators at these dims). Using 2 there also halves the template instantiations and simplifies the tail.
BF16/INT8/UINT8/SQ8↔SQ8 are intentionally untouched: BF16 AVX2 was benchmarked and showed no gain (its unpacklo/unpackhi halves already provide enough ILP), and the integer VNNI kernels have a different bottleneck profile.
Why (measured)
Standalone benchmark over the exact kernel templates (google-benchmark, pinned to a P-core, 9+ repetitions, random interleaving, medians; Intel Core Ultra 7 255U — AVX2/F16C class). The table below reports the multi-accumulator variant vs. the original single-accumulator kernel:
Newly-SIMD small dimensions vs. the scalar fallback they used before:
tests/benchmark/bm_spaces_*before/after on the same machine confirms (medians of 5 interleaved repetitions): FP32 SSE high-dim 1.38-1.48x, FP64 SSE high-dim 1.65-1.77x, FP16 F16C high-dim 1.76-1.82x, FP32/FP64 AVX 1.08-1.20x, residual dims improved across the board,_512_bit_chunksat dim 16/32 ~1.0x (no regression at the smallest sizes).AVX-512 (Xeon Platinum 8488C, via
bm-spacesCI): FP32/FP64 AVX-512 kernels ~1.2–1.35x over the previous single-accumulator versions. A dedicated 2-vs-4 comparison run showed 4 accumulators give no benefit on AVX-512 (once an ~11% inter-run instance offset — visible on every identical-code control family — is normalized out), which is why the AVX-512 kernels use 2. Full detail in the follow-up comment.Residual-load safety (why the new thresholds are safe)
_mm512_maskz_loadu_*), safe for any dimension.my_mm256_maskz_loadu_ps,_mm_loadu_si128+ blend). These read a full 8-element block from the start of the vector, so they are safe iff the vector spans at least 8 elements — hence thedim >= 8(FP32/FP16/SQ8) anddim >= 4(FP64, 4-double block) floors. Every subsequent full step stays withindimby construction (verified for all dims ≤ 4096 with a coverage simulation of every loop/tail structure).residual%4==3, which reads one aligned 16-byte block — in-bounds for every dim ≥ 8 that reaches it).Testing
test_spaces(1267 tests) passes locally, including the extended coverage added here: FP32 opt tests now run dims 8-32, FP64 4-16, FP16 8-64, SQ8-FP32 8-32, andsmallDimChooserasserts the new per-type scalar ranges.IP_AVX512F_SQ8_FP16.h, were exercised on thebm-spacesCI Xeon, and passed basic + sanitizer unit-test CI on the (now-closed) 2-vs-4 comparison branch [BENCH-ONLY] AVX-512: 2 accumulators vs 4 (do not merge) #986.main's spaces sources):*SerializationCurrentVersionMulti,HNSWSerializationV3(environmental), and atest_commonsegfault atSearchDifferentScoreswhen run after the rest of the suite.🤖 Generated with Claude Code
Note
Medium Risk
Touches hot-path distance kernels used in vector search; correctness depends on residual load bounds at small dimensions, though behavior is covered by expanded unit tests and documented safety rules.
Overview
x86 inner-product and L2 SIMD kernels (FP32/FP64/FP16 and SQ8↔FP32) now use 2–4 independent SIMD accumulators per loop instead of one, then fold once before reduction, to break FMA/mul→add dependency chains. Main loops switch from
do-whiletowhile, so they can run zero times when only the compile-time residual remains.Residual handling is reworked (
residual % lane_width, extra constexpr steps for 8/16/24-byte tails), template residual ranges widen (e.g. 0..31 on many FP32/SQ8 paths), andCHOOSE_IMPLEMENTATIONchunk sizes move to match 32-element (or 16 FP64 AVX-512 / 4 FP64 SSE) unrolling.Dispatchers in
IP_space.cpp/L2_space.cpplower SIMD cutoffs to what each kernel’s loads actually require (FP32/SQ8 from dim ≥ 8, FP64 from ≥ 4, FP16 tiered at 8 / 16 / 32 by ISA).test_spacesextends parameterized dim coverage and updatessmallDimChooserexpectations.Reviewed by Cursor Bugbot for commit 3b04e14. Bugbot is set up for automated code reviews on this repo. Configure here.